"""
Copyright (c) 2025 Huawei Technologies Co., Ltd.
openUBMC is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
"""

from dataclasses import dataclass, field
from typing import Dict, Any, List, TYPE_CHECKING

from ..base import IRNode


class Parameters:
    """注解参数"""

    def __init__(self):
        self.values: Dict[str, List[Any]] = {}

    def get(self, key: str, default: Any = None) -> Any:
        """获取参数的第一个值"""
        if not self.has(key):
            return default

        return self.values[key][0]

    def pop(self, key: str):
        """移除第一个参数值"""
        if not self.has(key):
            return

        self.values[key].pop(0)

    def add(self, key, value):
        """添加参数值"""
        if key not in self.values:
            self.values[key] = []
        self.values[key].append(value)

    def keys(self):
        """获取参数键列表"""
        return self.values.keys()

    def has(self, key: str) -> bool:
        """检查是否存在参数"""
        return (key in self.values) and (len(self.values[key]) > 0)

    def __str__(self) -> str:
        # 显示参数时，对于单值参数显示值，对于多值参数显示列表
        display_params = {}
        for key, values in self.values.items():
            if len(values) == 1:
                display_params[key] = values[0]
            else:
                display_params[key] = values
        return f"{display_params}"


@dataclass
class IRAnnotation(IRNode):
    """统一的注解 IR 节点

    通用的注解系统，支持任意注解名称和参数
    支持多值参数：parameters 存储为 Dict[str, List[Any]]
    """

    node_type: str = "annotation"

    # 注解基本信息
    annotation_name: str = field(default="")  # 注解名称，如 "mcpy_type_alias" 等
    parameters: Parameters = field(default_factory=Parameters)

    def get_parameter(self, key: str, default: Any = None) -> Any:
        """获取注解参数的第一个值"""
        return self.parameters.get(key, default)

    def get_parameters(self, key: str, default: Any = []) -> List[Any]:
        """获取注解参数的值列表"""
        if key not in self.parameters.values:
            return default
        return self.parameters.values[key]

    def pop_parameter(self, key: str):
        """移除第一个参数值"""
        self.parameters.pop(key)

    def has_parameter(self, key: str) -> bool:
        """检查是否存在参数"""
        return self.parameters.has(key)

    def keys(self):
        """获取参数键列表"""
        return self.parameters.keys()

    def first_key(self):
        if not self.parameters:
            return None
        return next(iter(self.parameters.keys()))

    def __str__(self) -> str:
        # 显示参数时，对于单值参数显示值，对于多值参数显示列表
        return f"IRAnnotation({self.annotation_name}, {self.parameters})"
